# -*- coding: utf-8 -*-
"""
网络设备自动化巡检脚本（会话保活版）
优化点：
1. 增强会话保活机制
2. 改进输出缓冲处理
3. 添加多级重试策略
"""

# ​********************​ 基础库导入 ​********************
import paramiko
import os
import time
import sys
import socket
from datetime import datetime
import ollama
import re  # 用于设备提示符匹配

# ​********************​ 配置区域 ​********************
SSH_PORT = 22          # SSH服务端口
SSH_USER = "sysadmin"  # 设备登录用户名
SSH_PASS = "Admin@h3c"# 设备登录密码
OLLAMA_MODEL = "deepseek-r1:14b"  # AI分析模型
OLLAMA_ENDPOINT = "http://localhost:11434"  # 本地模型服务地址
COMMAND_DELAY = 3      # 命令间隔时间（秒）
TIMEOUT = 30           # 连接和命令超时时间
MAX_RETRIES = 3        # 最大连接重试次数
CMD_RETRIES = 2        # 单命令最大重试次数
INSPECTION_COMMANDS = [  # 巡检命令列表
    'display version',
    'display cpu-usage',
    'display memory',
    'display current-configuration'
]
DEVICE_PROMPT = r'<[\w-]+>$'  # 设备提示符正则（匹配类似<H3C>的提示符）
# ​*************************************************

def get_timestamp():
    """生成标准化时间戳 YYYY-MM-DD HH:MM:SS 格式"""
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def collect_device_info(ip):
    """
    增强版设备信息采集（带会话保活和缓冲处理）
    优化点：
    - 三级重试机制（连接重试+命令重试+输出重试）
    - 动态等待机制替代固定延时
    - 缓冲区分块读取
    """
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    for retry in range(MAX_RETRIES):
        try:
            print(f"[{get_timestamp()}] 正在连接 {ip}...")
            # 建立带保活参数的连接[3,5](@ref)
            client.connect(ip, SSH_PORT, SSH_USER, SSH_PASS,
                          timeout=TIMEOUT,
                          allow_agent=False,    # 禁用密钥代理[5](@ref)
                          look_for_keys=False,  # 禁用密钥查找[3](@ref)
                          banner_timeout=30)    # 延长banner超时
            
            transport = client.get_transport()
            if transport:  # 添加空值校验
                transport.set_keepalive(15)  # 每15秒发送保活包
            
            # 创建交互式shell会话
            if not transport:  # 增加连接有效性校验
                raise paramiko.SSHException("SSH传输通道建立失败")
            shell = transport.open_session()
            shell.get_pty()  # 获取伪终端
            shell.invoke_shell()
            shell.settimeout(TIMEOUT)  # 设置socket超时
            
            output = f"\n=== {ip} 巡检结果 ===\n"
            
            for cmd in INSPECTION_COMMANDS:
                cmd_output = ""
                for cmd_retry in range(CMD_RETRIES):
                    try:
                        print(f"[{get_timestamp()}] 执行命令: {cmd}")
                        
                        # 发送命令（带换行符）
                        shell.send(f"{cmd}\n".encode('utf-8'))  # 转换为bytes类型
                        time.sleep(0.3)  # 增加初始缓冲等待
                        
                        # 动态等待输出
                        start_time = time.time()
                        while time.time() - start_time < TIMEOUT:
                            if shell.recv_ready():
                                chunk = shell.recv(4096).decode('utf-8', 'ignore')
                                cmd_output += chunk
                                # 检测设备提示符出现[3](@ref)
                                if re.search(DEVICE_PROMPT, cmd_output):
                                    break
                            time.sleep(0.5)  # 降低CPU占用
                        else:
                            raise paramiko.SSHException("命令响应超时")
                            
                        # 清理命令回显
                        cleaned_output = re.sub(r'^.*?\n', '', cmd_output, flags=re.DOTALL)
                        output += f"[{cmd}]\n{cleaned_output.strip()}\n\n"
                        break  # 命令执行成功则跳出重试循环
                        
                    except paramiko.SSHException as e:
                        print(f"[{cmd}] 命令执行异常: {str(e)}")
                        if cmd_retry == CMD_RETRIES - 1:
                            output += f"[{cmd}] 错误：超过最大重试次数\n\n"
                        else:
                            time.sleep(2)  # 命令重试间隔
            
            return output

        except (paramiko.AuthenticationException, paramiko.SSHException) as e:
            print(f"[{ip}] 连接异常: {str(e)}")
            if retry == MAX_RETRIES - 1:
                return f"\n=== {ip} 连接失败 ===\n错误类型: {type(e).__name__}\n信息: {str(e)}\n"
            time.sleep(5)  # 指数退避等待
            
        finally:
            try:
                client.close()  # 确保连接关闭[2](@ref)
            except:
                pass

def generate_ai_report(content):
    """
    Ollama 大模型报告生成函数
    :param content: 原始巡检数据
    :return: 结构化分析报告（含错误处理）
    """
    try:
        print(f"[{get_timestamp()}] 正在生成分析报告...")
        response = ollama.chat(
            model=OLLAMA_MODEL,
            messages=[
                {"role": "system", "content": "你是一个专业的网络设备分析专家，能够根据设备巡检日志生成分析报告"},
                {"role": "user", "content": f"请根据以下设备巡检日志生成分析报告：\n{content}"}
            ]
        )
        return response['message']['content'].strip()
    except Exception as e:
        return f"Ollama 服务异常：{str(e)}"

def check_ollama_service():
    """
    Ollama 服务状态检测函数
    确保本地 Ollama 服务已正常启动
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(3)
            s.connect(('localhost', 11434))
    except Exception as e:
        print(f"[{get_timestamp()}] 错误：Ollama 服务未启动！请执行：ollama run {OLLAMA_MODEL}")
        sys.exit(1)

def main():
    """主执行流程"""
    # 生成带时间戳的报告文件
    report_file = f"Analysis_Report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
    
    try:
        # 读取设备IP列表
        with open('ip.txt', 'r') as f:
            ips = [line.strip() for line in f if line.strip()]
        
        if not ips:
            print(f"[{get_timestamp()}] 错误：ip.txt中没有找到有效IP地址")
            return

        # 执行设备巡检
        report_list = []  # 使用列表代替字符串拼接
        for ip in ips:
            print(f"\n[{get_timestamp()}] 开始巡检 {ip}")
            report_list.append(collect_device_info(ip))
        
        full_report = ''.join(report_list)  # 最后合并字符串

        # 生成中间文件
        with open('temp.txt', 'w') as f:
            f.write(full_report)

        # 调用大模型分析
        with open('temp.txt', 'r') as f:
            ai_report = generate_ai_report(f.read())

        # 保存最终报告
        with open(report_file, 'w') as f:
            f.write(f"=== 设备巡检分析报告 ===\n")
            f.write(f"生成时间: {get_timestamp()}\n")
            f.write(f"分析模型: {OLLAMA_MODEL}\n")
            f.write(f"巡检设备: {len(ips)}台\n\n")
            f.write(ai_report)

        print(f"\n[{get_timestamp()}] 报告生成成功: {report_file}")

    except FileNotFoundError:
        print(f"[{get_timestamp()}] 错误：未找到ip.txt文件")
    except Exception as e:
        print(f"[{get_timestamp()}] 发生未预期错误: {str(e)}")
    finally:
        # 清理临时文件
        time.sleep(10)
        if os.path.exists('temp.txt'):
            os.remove('temp.txt')
            print(f"[{get_timestamp()}] 已清理临时文件")

if __name__ == "__main__":
    check_ollama_service()
    main()